#!/home/craigb/bin/python

import os
import sys

from optparse import OptionParser

import tables


op = OptionParser(usage="Usage: %prog [opts] file1.h5 [file2.h5 [...]]")

op.add_option("-o", "--output", dest='output', default=None,
              help='Redirect output to FILE', metavar='FILE')
op.add_option("-n", "--nt", dest="nt_name", default="nt",
              help="Print length of table named `NT'", metavar="NT")
op.add_option("--file-list", dest="file_list", default='', 
              help="Use files listed in PATH for input", metavar="PATH")

op.add_option("-s", "--sum", dest='sum', action='store_true', default=False,
              help="Print only the sum total of lengths of all files.")

opts, args = op.parse_args()

if len(args) == 0:
    if opts.file_list:
        fp = open(opts.file_list, 'r')
        new_args = [line.strip() for line in fp]
        fp.close()
        
        new_args.sort()
        args = new_args

    else:
        sys.stderr.write("Please supply a node name and at least one HDF5 input file.\n")
        sys.stderr.write("See `pytlen.py -h' for more.\n")
        sys.stderr.flush()
        sys.exit(-1)

tot_length = 0

if opts.output:
    fpout = open(opts.output, 'w')
    fpout.write('{\n')

for fname in args:
    if not tables.isPyTablesFile(fname):
        sys.stderr.write("File `{0}' is not a valid PyTables file\n".format(fname))
        sys.stderr.write("Continuing...\n");
        sys.stderr.flush()
        continue
        
    
    try:
        fp = tables.openFile(fname, 'r')
    except IOError, e:
        sys.stderr.write("Error opening file `{0}':\n".format(fname))
        sys.stderr.write(str(e) + '\n')
        sys.stderr.write("Continuing...\n");
        sys.stderr.flush()
        continue
        
    try:
        length = eval("len(fp.root." + opts.nt_name + ")")
    except tables.NoSuchNodeError:
        sys.stderr.write("No node `fp.root." + opts.nt_name + "' in file " + fname + "\n")
        sys.stderr.write("Continuing...\n");
        sys.stderr.flush()
        fp.close()
        continue
    except AttributeError:
        sys.stderr.write("Root node not present in file " + fname + "\n")
        sys.stderr.write("Continuing...\n");
        sys.stderr.flush()
        fp.close()
        continue

    fp.close()
    if opts.output:
        fpout.write("'{0}': {1},".format(os.path.basename(fname), length))
        fpout.write('}\n')
        fpout.close()
    elif not opts.sum:
        print "{0} {1}".format(fname, length)
    tot_length += length

print "# of files examined: {0}; total length is {1}".format(len(args), tot_length)

